import backtrader as bt
import datetime
import pandas as pd

class MyStrategy(bt.Strategy):
    params = (
        ("sma_period", 20),
    )

    def __init__(self):
        self.dataclose = self.data.close
        self.sma = bt.indicators.SimpleMovingAverage(self.data.close, period=self.params.sma_period)
        self.order = None

    def next(self):
        if self.dataclose > self.sma:
            # 如果收盘价超过移动平均线，发出买入信号
            if self.order is None:
                self.order = self.buy()
        elif self.dataclose < self.sma:
            # 如果收盘价低于移动平均线，发出卖出信号
            if self.order is None:
                self.order = self.sell()

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"买入: {order.executed.price}")
            elif order.issell():
                self.log(f"卖出: {order.executed.price}")

            self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f"交易利润: {trade.pnl}")

if __name__ == "__main__":
    cerebro = bt.Cerebro()

    # 加载实时数据（模拟实时数据）
    data = bt.feeds.PandasData(dataname=pd.DataFrame(columns=["Open", "High", "Low", "Close", "Volume"]),
                               historical=False)

    cerebro.adddata(data)

    # 添加策略
    cerebro.addstrategy(MyStrategy)

    # 设置初始资金
    cerebro.broker.set_cash(100000)

    # 设置手续费
    cerebro.broker.setcommission(commission=0.001)

    print(f"初始资金: {cerebro.broker.getvalue()}")

    # 运行回测
    cerebro.run()

    print(f"结束资金: {cerebro.broker.getvalue()}")
